#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MCP 服务器核心模块

实现 MySQL Model Context Protocol 服务器，
提供数据库查询、表结构获取等功能。
"""

import time
from typing import Any, Dict, List, Optional, Sequence
from mcp.server import Server
from mcp.server.models import InitializationOptions
from mcp.server.stdio import stdio_server
from mcp.types import (
    Resource,
    Tool,
    TextContent,
    ImageContent,
    EmbeddedResource,
    LoggingLevel
)
import asyncio

from .config import Config, get_config
from .database import DatabaseManager
from .logger import (
    setup_logger,
    get_logger,
    log_mcp_server_start,
    log_mcp_server_stop,
    log_mcp_request,
    log_mcp_response,
    log_config_loaded
)
from .utils import (
    validate_sql_query,
    serialize_mysql_results,
    format_error_message,
    safe_json_dumps,
    sanitize_sql_identifier,
    validate_table_name
)


class MySQLMCPServer:
    """
    MySQL MCP 服务器类
    
    实现 Model Context Protocol 服务器，提供 MySQL 数据库访问功能。
    支持自定义端口配置和完善的错误处理。
    """
    
    def __init__(self, config: Optional[Config] = None):
        """
        初始化 MCP 服务器
        
        Args:
            config: 配置对象，如果为 None 则使用默认配置
        """
        self.config = config or get_config()
        
        # 设置日志
        self.std_logger = setup_logger(self.config.log)
        self.logger = get_logger("mcp_server")
        
        # 记录配置加载
        log_config_loaded(self.config.to_dict())
        
        # 初始化数据库管理器
        self.db_manager = DatabaseManager(self.config.database)
        
        # 创建 MCP 服务器
        self.server = Server(self.config.mcp.server_name)
        
        # 注册处理器
        self._register_handlers()
    
    def _register_handlers(self) -> None:
        """
        注册 MCP 服务器处理器
        """
        # 注册资源处理器
        @self.server.list_resources()
        async def list_resources() -> List[Resource]:
            """
            列出可用资源
            
            Returns:
                List[Resource]: 资源列表
            """
            return [
                Resource(
                    uri="mysql://database/info",
                    name="数据库信息",
                    description="获取数据库基本信息和表列表",
                    mimeType="application/json"
                ),
                Resource(
                    uri="mysql://database/tables",
                    name="数据库表列表",
                    description="获取数据库中所有表的列表",
                    mimeType="application/json"
                )
            ]
        
        @self.server.read_resource()
        async def read_resource(uri: str) -> str:
            """
            读取资源内容
            
            Args:
                uri: 资源 URI
                
            Returns:
                str: 资源内容
                
            Raises:
                ValueError: 当资源不存在时
            """
            start_time = time.time()
            
            try:
                log_mcp_request("read_resource", {"uri": uri})
                
                if uri == "mysql://database/info":
                    info = self.db_manager.get_database_info()
                    result = safe_json_dumps(info, indent=2)
                elif uri == "mysql://database/tables":
                    info = self.db_manager.get_database_info()
                    tables = info.get('tables', [])
                    result = safe_json_dumps(tables, indent=2)
                else:
                    raise ValueError(f"未知的资源 URI: {uri}")
                
                response_time = time.time() - start_time
                log_mcp_response("read_resource", True, response_time)
                
                return result
                
            except Exception as e:
                response_time = time.time() - start_time
                log_mcp_response("read_resource", False, response_time)
                raise ValueError(format_error_message(e, "读取资源"))
        
        # 注册工具处理器
        @self.server.list_tools()
        async def list_tools() -> List[Tool]:
            """
            列出可用工具
            
            Returns:
                List[Tool]: 工具列表
            """
            return [
                Tool(
                    name="execute_query",
                    description="执行 SQL 查询语句",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "query": {
                                "type": "string",
                                "description": "要执行的 SQL 查询语句（仅支持 SELECT）"
                            },
                            "limit": {
                                "type": "integer",
                                "description": "结果限制数量（可选）",
                                "minimum": 1,
                                "maximum": 1000,
                                "default": 100
                            }
                        },
                        "required": ["query"]
                    }
                ),
                Tool(
                    name="get_table_schema",
                    description="获取表结构信息",
                    inputSchema={
                        "type": "object",
                        "properties": {
                            "table_name": {
                                "type": "string",
                                "description": "表名称"
                            }
                        },
                        "required": ["table_name"]
                    }
                ),
                Tool(
                    name="list_tables",
                    description="列出数据库中的所有表",
                    inputSchema={
                        "type": "object",
                        "properties": {},
                        "required": []
                    }
                ),
                Tool(
                    name="get_database_info",
                    description="获取数据库基本信息",
                    inputSchema={
                        "type": "object",
                        "properties": {},
                        "required": []
                    }
                )
            ]
        
        @self.server.call_tool()
        async def call_tool(name: str, arguments: Dict[str, Any]) -> List[TextContent]:
            """
            调用工具
            
            Args:
                name: 工具名称
                arguments: 工具参数
                
            Returns:
                List[TextContent]: 工具执行结果
                
            Raises:
                ValueError: 当工具不存在或执行失败时
            """
            start_time = time.time()
            
            try:
                log_mcp_request(name, arguments)
                
                if name == "execute_query":
                    result = await self._execute_query_tool(arguments)
                elif name == "get_table_schema":
                    result = await self._get_table_schema_tool(arguments)
                elif name == "list_tables":
                    result = await self._list_tables_tool(arguments)
                elif name == "get_database_info":
                    result = await self._get_database_info_tool(arguments)
                else:
                    raise ValueError(f"未知的工具: {name}")
                
                response_time = time.time() - start_time
                log_mcp_response(name, True, response_time)
                
                return [TextContent(type="text", text=result)]
                
            except Exception as e:
                response_time = time.time() - start_time
                log_mcp_response(name, False, response_time)
                error_msg = format_error_message(e, f"工具 {name}")
                return [TextContent(type="text", text=f"错误: {error_msg}")]
    
    async def _execute_query_tool(self, arguments: Dict[str, Any]) -> str:
        """
        执行查询工具
        
        Args:
            arguments: 工具参数
            
        Returns:
            str: 查询结果 JSON 字符串
            
        Raises:
            ValueError: 当查询无效或执行失败时
        """
        query = arguments.get("query", "").strip()
        limit = arguments.get("limit", 100)
        
        if not query:
            raise ValueError("查询语句不能为空")
        
        # 验证查询
        validate_sql_query(query)
        
        # 添加 LIMIT 子句（如果查询中没有）
        query_upper = query.upper()
        if "LIMIT" not in query_upper:
            query += f" LIMIT {limit}"
        
        # 执行查询
        results = self.db_manager.execute_query(query)
        
        # 序列化结果
        serialized_results = serialize_mysql_results(results)
        
        return safe_json_dumps({
            "query": query,
            "row_count": len(serialized_results),
            "results": serialized_results
        }, indent=2)
    
    async def _get_table_schema_tool(self, arguments: Dict[str, Any]) -> str:
        """
        获取表结构工具
        
        Args:
            arguments: 工具参数
            
        Returns:
            str: 表结构信息 JSON 字符串
            
        Raises:
            ValueError: 当表名无效或获取失败时
        """
        table_name = arguments.get("table_name", "").strip()
        
        if not table_name:
            raise ValueError("表名称不能为空")
        
        # 验证表名
        validate_table_name(table_name)
        
        # 获取表结构
        table_info = self.db_manager.get_table_info(table_name)
        
        return safe_json_dumps(table_info, indent=2)
    
    async def _list_tables_tool(self, arguments: Dict[str, Any]) -> str:
        """
        列出表工具
        
        Args:
            arguments: 工具参数
            
        Returns:
            str: 表列表 JSON 字符串
        """
        # 获取数据库信息
        db_info = self.db_manager.get_database_info()
        tables = db_info.get('tables', [])
        
        return safe_json_dumps({
            "database": db_info.get('database_name'),
            "table_count": len(tables),
            "tables": tables
        }, indent=2)
    
    async def _get_database_info_tool(self, arguments: Dict[str, Any]) -> str:
        """
        获取数据库信息工具
        
        Args:
            arguments: 工具参数
            
        Returns:
            str: 数据库信息 JSON 字符串
        """
        # 获取数据库信息
        db_info = self.db_manager.get_database_info()
        
        return safe_json_dumps(db_info, indent=2)
    
    async def start(self) -> None:
        """
        启动 MCP 服务器
        
        Raises:
            Exception: 当启动失败时
        """
        try:
            # 连接数据库
            self.db_manager.connect()
            
            # 记录服务器启动
            log_mcp_server_start(
                self.config.mcp.server_name,
                self.config.mcp.server_version
            )
            
            # 启动服务器
            async with stdio_server() as (read_stream, write_stream):
                await self.server.run(
                    read_stream,
                    write_stream,
                    InitializationOptions(
                        server_name=self.config.mcp.server_name,
                        server_version=self.config.mcp.server_version,
                        capabilities=self.server.get_capabilities()
                    )
                )
                
        except Exception as e:
            self.logger.error("MCP 服务器启动失败", error=str(e), exc_info=True)
            raise
        finally:
            await self.stop()
    
    async def stop(self) -> None:
        """
        停止 MCP 服务器
        """
        try:
            # 断开数据库连接
            if self.db_manager.is_connected:
                self.db_manager.disconnect()
            
            # 记录服务器停止
            log_mcp_server_stop()
            
        except Exception as e:
            self.logger.error("停止 MCP 服务器时发生错误", error=str(e))
    
    def run(self) -> None:
        """
        运行 MCP 服务器（同步方法）
        """
        try:
            asyncio.run(self.start())
        except KeyboardInterrupt:
            self.logger.info("收到中断信号，正在停止服务器...")
        except Exception as e:
            self.logger.error("服务器运行时发生错误", error=str(e), exc_info=True)
            raise


async def run_server(config: Optional[Config] = None) -> None:
    """
    运行 MCP 服务器（异步函数）
    
    Args:
        config: 配置对象
    """
    server = MySQLMCPServer(config)
    await server.start()


def main() -> None:
    """
    主函数
    
    用于命令行启动服务器。
    """
    try:
        # 加载配置
        config = get_config()
        
        # 验证配置
        config.validate()
        
        # 创建并运行服务器
        server = MySQLMCPServer(config)
        server.run()
        
    except KeyboardInterrupt:
        print("\n服务器已停止")
    except Exception as e:
        print(f"启动服务器失败: {e}")
        raise


if __name__ == "__main__":
    main()